import json
from pyparsing import *

class ModelParser():
    def __init__(self):
        self._setup_parser()

    def parse(self, text):
        try:
            parsed_data = self.parser.parseString(text.strip(), parseAll=True)
            return self._extract_data(parsed_data)
        except ParseException as pe:
            print(f"Failed to parse at line {pe.lineno} col {pe.column}: {pe}")
            self._highlight_error_position(text, pe.lineno, pe.column)
            return None

    def _highlight_error_position(self, text, lineno, column):
        # 打印错误文本的上下文标识不匹配位置
        lines = text.splitlines()
        if lineno <= len(lines):
            error_line = lines[lineno - 1]
            print(f"{error_line}")
            print(" " * (column - 1) + "^")
            print("")
        else:
            print('Error location is beyond the input text length.')

    def _setup_parser(self):
        # 基本元素
        integer = Word(nums).setParseAction(lambda t: int(t[0]))
        floatnum = Regex(r'[+-]?\d+(\.\d*)?([eE][+-]?\d+)?').setParseAction(lambda t: float(t[0]))
        identifier = Word(alphas, alphanums + "_") | Word(alphanums + "-" + alphanums) | integer
        equals = Literal("=")
        bool_none = oneOf("True False None").setParseAction(lambda t: t[0] == "True")
        multiply = integer + Optional(White()) + Literal("x")

        # 字符串解析器
        string_literal = quotedString.setParseAction(removeQuotes)

        # 定义Layer的参数
        simple_value = floatnum | integer | string_literal | self._tuple_parser() | bool_none | identifier
        arg = Group(identifier + equals + simple_value) | integer
        arguments = delimitedList(arg, ',')

        # 层解析，包括嵌套结构
        value = Forward()

        # 普通的layer层定义
        layer_base = Group(identifier.setResultsName("layer_base_name") + Suppress("(") + Optional(arguments.setResultsName("layer_base_args")) + Suppress(")")).setResultsName("layer_base")

        # 普通的block_entry
        normal_entry = Group(Suppress("(") + identifier.setResultsName("block_entry_name") + Suppress("):") + ZeroOrMore(multiply).setResultsName("multiply_entries") + value.setResultsName("block_entry_value")).setResultsName("block_entry")

        # 定义init_cfg开头的entry
        cfg_entry = Group(Literal("init_cfg").setResultsName("block_entry_name") + equals + SkipTo(LineEnd()).setResultsName("block_entry_value"))
        
        # 定义value为None的特殊entry
        none_value = Group(ZeroOrMore(integer+multiply).setResultsName("multiply_entries") + bool_none).setResultsName("block_entry_value")
        none_entry = Group(Suppress("(") + identifier.setResultsName("block_entry_name") + Suppress("):") + ZeroOrMore(multiply).setResultsName("multiply_entries") + none_value.setResultsName("block_entry_value")).setResultsName("block_entry")

        # 定义递归的block_entry类型 
        block_entry = none_entry | normal_entry | cfg_entry
        layer_block = Group(identifier.setResultsName("layer_block_name") + Suppress("(") + ZeroOrMore(block_entry).setResultsName("block_entries") + Suppress(")")).setResultsName("layer_block")

        # value可以是layer_base或者嵌套的layer_block
        value <<= layer_base | layer_block

        # 定义解析器
        self.parser = layer_block

    def _tuple_parser(self):
        lparen = Suppress("(")
        rparen = Suppress(")")
        integer = Word(nums).setParseAction(lambda t: int(t[0]))
        return Group(lparen + delimitedList(integer) + rparen)
    
    def _extract_data(self, parsed_data):
        data = self._parse_data( parsed_data[0] )
        return data
    
    def _parse_data(self, data):
        properties = [ 'layer_block',  'block_entry', 'layer_base' ]
        if isinstance(data, ParseResults):
            return self._parse_layer_block( data )
        else:
            print('Unknown parsed_data = ', data)

    def _parse_layer_block(self, layer_block):
        json_layer_block = {
            'block_name': layer_block.layer_block_name,
            'block_entries': []
        }
        for block_entry in layer_block.block_entries:
            json_block = self._parse_block_entry( block_entry )
            json_layer_block['block_entries'].append( json_block )
        return json_layer_block
    
    def _parse_block_entry(self, block_entry):
        json_block_entry = {
            'entry_name': block_entry.block_entry_name,
            'entry_times': block_entry.multiply_entries[0] if len(block_entry.multiply_entries) > 0 else 1,
            'entry_value': None
        }
        block_entry_value = block_entry.block_entry_value
        if( block_entry.block_entry_name == 'init_cfg' ):
            json_block_entry['entry_value'] = self._parse_init_cfg( block_entry_value )
        else:
            if 'layer_block_name' in block_entry_value:
                json_block_entry['entry_value'] = self._parse_layer_block( block_entry_value ) 
            elif 'layer_base_name' in block_entry_value:
                json_block_entry['entry_value'] = self._parse_layer_base( block_entry_value )
        return json_block_entry
    
    def _parse_layer_base(self, layer_base):
        json_layer_base = {
            'layer_name' : layer_base.layer_base_name,
            'layer_args': []
        }
        for arg in layer_base.layer_base_args:
            if isinstance(arg, int):
                json_layer_base['layer_args'].append( arg )
            else:
                arg_name, arg_value = arg[0], tuple(arg[2]) if isinstance(arg[2], ParseResults) else arg[2]
                json_layer_base['layer_args'].append( { arg_name : arg_value } )
        return json_layer_base
    
    def _parse_init_cfg(self, init_cfg_value):
        # print(init_cfg_value)
        json_init_cfg = json.loads(init_cfg_value.replace("'", '"'))
        return json_init_cfg
    
    def to_json(self, data):
        if data is None:
            return json.dumps({})
        return json.dumps(data, indent=2)

if __name__ == "__main__":
    data = ""
    with open('./model_arch.txt', 'rt') as fd:
        data = fd.read()
    # print(data)

    model_parser = ModelParser()
    model_info = model_parser.parse(data)
    print(model_parser.to_json(model_info))
    # print(model_parser.to_json(model_info))
